import numpy as np  
import faiss  
  
class FaissKNeighbors:  
    def __init__(self, k: int = 1, res: faiss.GpuResources = None):  
        self.index = None  # FAISS索引，用于存储训练数据  
        self.y = None  # 训练数据的标签  
        self.k = k  # 最近邻个数  
        self.res = res  # FAISS GPU资源对象  
  
    def fit(self, X: np.ndarray, y: np.ndarray):  
        # 检查输入数据类型和形状  
        if not isinstance(X, np.ndarray) or not isinstance(y, np.ndarray):  
            raise TypeError("X and y must be numpy arrays.")  
        if X.shape[0] != y.shape[0]:  
            raise ValueError("X and y must have the same number of samples.")  
  
        # 使用内积索引（假设数据已归一化）  
        self.index = faiss.IndexFlatIP(X.shape[1])  
  
        # 如果提供了GPU资源，则将索引转移到GPU上  
        if self.res is not None:  
            if faiss.gpu_available():  
                self.index = faiss.index_cpu_to_gpu(self.res, 0, self.index)  
            else:  
                print("GPU not available, using CPU instead.")  
  
        self.index.add(X.astype(np.float32))  
        self.y = y  
  
    def predict(self, X: np.ndarray) -> np.ndarray:  
        distances, indices = self.index.search(X.astype(np.float32), self.k)  
          
        # 使用字典来存储投票，避免重复计算  
        vote_dict = {}  
        for i, indices_i in enumerate(indices):  
            for idx, label in zip(indices_i, self.y[indices_i]):  
                if label not in vote_dict:  
                    vote_dict[label] = 0  
                vote_dict[label] += 1  
  
        # 获取预测标签  
        predictions = np.array([max(vote_dict, key=vote_dict.get)[0] for _ in range(X.shape[0])])  
        return predictions  
  
    def score(self, X: np.ndarray, y: np.ndarray) -> float:  
        predictions = self.predict(X)  
        accuracy = np.mean(predictions == y)  
        return accuracy
